import numpy as np
from utils import *

class BernoulliBandit:
    '''伯努利多臂老虎机，输入K表示拉杆个数'''
    def __init__(self,K):
        # 随机生成k个0~1的数，作为拉动每根拉杆的获奖概率
        self.probs = np.random.uniform(size=K)
        # 获奖概率最大的拉杆索引
        self.best_idx = np.argmax(self.probs)
        # 最大的获奖概率
        self.best_prob = self.probs[self.best_idx]
        self.K = K

    def step(self,k):
        # 当玩家选择了k号拉杆后，根据拉动该老虎机的k号拉杆得到奖励的概率返回1（获奖）或0（未获奖）
        if np.random.rand() < self.probs[k]:
            return 1
        else:
            return 0

class Solver:
    """ 多臂老虎机算法基本框架 """
    def __init__(self, bandit):
        self.bandit = bandit
        self.counts = np.zeros(self.bandit.K) # 每根拉杆的尝试次数
        self.regret = 0. # 当前步的累积懊悔
        self.actions = [] # 维护一个列表,记录每一步的动作
        self.regrets = [] # 维护一个列表,记录每一步的累积懊悔
    def update_regret(self, k):
        # 计算累积懊悔并保存,k为本次动作选择的拉杆的编号
        self.regret += self.bandit.best_prob - self.bandit.probs[k]
        self.regrets.append(self.regret)
    def run_one_step(self):
        # 返回当前动作选择哪一根拉杆,由每个具体的策略实现
        # 根据策略选择动作、根据动作获取奖励、更新期望奖励估值
        raise NotImplementedError
    def run(self, num_steps):
        # 运行一定次数,num_steps为总运行次数
        for _ in range(num_steps):
            k = self.run_one_step()
            self.counts[k] += 1     # 4.更新计数
            self.actions.append(k)
            self.update_regret(k)   # 5.更新累积懊悔

class EpsilonGreedy(Solver):
    """ epsilon 贪婪算法,继承 Solver 类 """
    def __init__(self, bandit, epsilon=0.01, init_prob=1.0):
        super(EpsilonGreedy, self).__init__(bandit)
        self.epsilon = epsilon
        #初始化拉动所有拉杆的期望奖励估值
        self.estimates = np.array([init_prob] * self.bandit.K)
    def run_one_step(self):
        # 1.根据策略选择动作
        if np.random.random() < self.epsilon:
            k = np.random.randint(0, self.bandit.K)  # 随机选择一根拉杆
        else:
            k = np.argmax(self.estimates)  # 选择期望奖励估值最大的拉杆
        # 2.根据动作获取奖励
        r = self.bandit.step(k)  # 得到本次动作的奖励
        # 3.更新期望奖励估值
        self.estimates[k] += 1. / (self.counts[k] + 1) * (r - self.estimates[k])
        return k

class DecayingEpsilonGreedy(Solver):
    '''epsilon值随时间衰减的epsilon-贪婪算法'''
    def __init__(self,bandit,init_prob=1.0):
        super(DecayingEpsilonGreedy,self).__init__(bandit)
        self.estimates = np.array([init_prob] * self.bandit.K)
        self.total_count = 0

    def run_one_step(self):
        self.total_count += 1
        if np.random.random() < 1/self.total_count:  # epsilon值随时间衰减
            k = np.random.randint(0,self.bandit.K)
        else:
            k = np.argmax(self.estimates)

        r = self.bandit.step(k)
        self.estimates[k] += 1. / (self.counts[k]+1) * (r-self.estimates[k])
        return k

class UCB(Solver):
    '''UCB（上置信界）算法'''
    def __init__(self,bandit,coef,init_prob=1.0):
        super(UCB,self).__init__(bandit)
        self.coef = coef
        self.estimates = np.array([init_prob] * self.bandit.K)
        self.total_count = 0

    def run_one_step(self):
        self.total_count += 1
        unsure = np.sqrt(np.log(self.total_count) / (2*(self.counts+1)))  # 计算不确定性度量
        ucb = self.estimates + self.coef * unsure  # 计算上置信界
        k = np.argmax(ucb)  # 选出上置信界最大的拉杆
        r = self.bandit.step(k)
        self.estimates[k] += 1. / (self.counts[k]+1) * (r-self.estimates[k])
        return k

class ThomsonSampling(Solver):
    """汤普森采样算法"""
    def __init__(self,bandit):
        super(ThomsonSampling,self).__init__(bandit)
        self._a = np.ones(self.bandit.K)  # 列表，表示每根拉杆奖励为1的次数
        self._b = np.ones(self.bandit.K)  # 列表，表示每根拉杆奖励为0的次数

    def run_one_step(self):
        samples = np.random.beta(self._a,self._b)  # 按照beta分布采样一组奖励样本
        k = np.argmax(samples)  # 选出采样奖励最大的拉杆
        r = self.bandit.step(k)

        self._a[k] += r  # 更新Beta分布的第一个参数
        self._b[k] += 1-r  # 更新Beta分布的第二个参数
        return k

if __name__ == '__main__':
    np.random.seed(0)
    K = 10
    bandit_10_arm = BernoulliBandit(K)
    print(f"随机生成了一个{K}臂老虎机")
    print(f"获奖概率最大的拉杆为{bandit_10_arm.best_idx}号,其获奖概率为{bandit_10_arm.best_prob:.4f}")

    epsilon_greedy_solver = EpsilonGreedy(bandit_10_arm,epsilon=0.01)
    epsilon_greedy_solver.run(5000)
    print('epsilon-贪婪算法的累积懊悔为： ',epsilon_greedy_solver.regret)
    plot_results([epsilon_greedy_solver],["EpsilonGreedy"])

    epsilons = [1e-4,0.01,0.1,0.25,0.5]
    epsilon_greedy_solver_list = [EpsilonGreedy(bandit_10_arm,epsilon=e) for e in epsilons]
    epsilon_greedy_names = [f"epsilon={e}" for e in epsilons]
    for solver in epsilon_greedy_solver_list:
        solver.run(5000)
    plot_results(epsilon_greedy_solver_list,epsilon_greedy_names)

    decaying_epsilon_greedy_solver = DecayingEpsilonGreedy(bandit_10_arm)
    decaying_epsilon_greedy_solver.run(5000)
    print('epsilon值随时间衰减的epsilon-贪婪算法的累积懊悔为： ',decaying_epsilon_greedy_solver.regret)
    plot_results([decaying_epsilon_greedy_solver],["DecayingEpsilonGreedy"])

    coef = 1  # 控制不确定性比重的系数
    UCB_Solver = UCB(bandit_10_arm,coef)
    UCB_Solver.run(5000)
    print('UCB算法的累积懊悔为： ',UCB_Solver.regret)
    plot_results([UCB_Solver],["UCB"])

    thompson_sampling_solver = ThomsonSampling(bandit_10_arm)
    thompson_sampling_solver.run(5000)
    print('汤普森采样算法的累积懊悔为： ',thompson_sampling_solver.regret)
    plot_results([thompson_sampling_solver],["ThompsonSampling"])
